/**
 * LY.com Inc.
 * Copyright (c) 2004-2021 All Rights Reserved.
 */
package consistenthashing;

import java.util.*;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * @author bixiao.wang
 * @version $Id: ConsistentHashing, v 0.1 2021/9/17 16:50 bixiao.wang Exp $
 */
public abstract class AbstractConsistentHashing {

    /**
     * 每个节点的虚拟节点数量，0和1表示不使用虚拟节点
     * 设置后不支持修改
     */
    private int virtualNodeNum;
    private boolean useVirtualNode;

    Map<String, List<Integer>> nodeMap = new HashMap<>();
    TreeMap<Integer, String> hashingMap = new TreeMap<>();

    private ReadWriteLock readWriteLock = new ReentrantReadWriteLock();


    public AbstractConsistentHashing(int virtualNodeNum) {
        this.virtualNodeNum = virtualNodeNum;
        this.useVirtualNode = virtualNodeNum > 1;
    }

    /**
     * 添加节点
     * @param node
     */
    public void addNode(final String node) {
        final Lock lock = readWriteLock.writeLock();
        try {
            lock.lock();
            if (useVirtualNode) {
                // 使用虚拟节点
                List<Integer> virtualList = new ArrayList<>(virtualNodeNum);
                nodeMap.put(node, virtualList);
                Random random = new Random();
                for (int i = 0; i < virtualNodeNum; i++) {
                    int slot = random.nextInt();
                    while (hashingMap.containsKey(slot)) {
                        // 哈希冲突重新计算hash
                        slot = random.nextInt();
                    }
                    virtualList.add(slot);
                    hashingMap.put(slot, node);
                }
            } else {
                // 不使用虚拟节点
                Integer slot = hashing(node);
                while (hashingMap.containsKey(slot)) {
                    // 哈希冲突重新计算hash
                    slot = hashing(node + slot);
                }
                hashingMap.put(slot, node);
                nodeMap.put(node, Arrays.asList(slot));
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 移除节点
     * @param node
     */
    public void removeNode(final String node) {
        final Lock lock = readWriteLock.writeLock();
        try {
            lock.lock();
            List<Integer> slots = nodeMap.remove(node);
            for (Integer slot : slots) {
                hashingMap.remove(slot);
            }
        } finally {
            lock.unlock();
        }
    }

    public String allocateNode(final String key) {
        final Lock lock = readWriteLock.readLock();
        try {
            lock.lock();
            int hashing = hashing(key);
            Map.Entry<Integer, String> higherEntry = hashingMap.higherEntry(hashing);
            if (higherEntry == null) {
                Map.Entry<Integer, String> firstEntry = hashingMap.firstEntry();
                return firstEntry == null ? null : firstEntry.getValue();
            } else {
                return higherEntry.getValue();
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 哈希算法，越散列效果越好
     * @param node
     * @return
     */
    public abstract int hashing(String node);
}
